/*
 * Copyright (c) 2024 OpenHiTLS. All rights reserved.
 * Licensed under the OpenHiTLS license.
 */
package org.openhitls.tls;

import org.junit.BeforeClass;
import org.junit.Test;
import static org.junit.Assert.*;

import javax.net.ssl.*;
import java.io.File;
import java.nio.ByteBuffer;
import java.security.*;
import java.security.cert.X509Certificate;

import org.openhitls.crypto.jce.provider.HiTls4jProvider;

/**
 * Tests for SSLEngine handshake and data transfer operations.
 * These tests simulate client-server TLS communication using SSLEngine.
 */
public class SSLEngineHandshakeTest {

    private static final String PROVIDER_NAME = HiTls4jProvider.PROVIDER_NAME;

    // Buffer sizes based on TLS record limits
    private static final int APP_BUFFER_SIZE = 16384;
    private static final int NET_BUFFER_SIZE = 17408; // TLS record max + header

    @BeforeClass
    public static void loadNativeLibraries() {
        String libraryPath = new File(System.getProperty("user.dir") + "/target/native/libhitls_crypto_jni.so").getAbsolutePath();
        System.load(libraryPath);
        Security.addProvider(new HiTls4jProvider());
    }

    /**
     * Creates a trust-all SSLContext for testing purposes.
     */
    private SSLContext createTestSSLContext() throws Exception {
        SSLContext ctx = SSLContext.getInstance("TLS", PROVIDER_NAME);

        TrustManager[] trustAllCerts = new TrustManager[] {
            new X509TrustManager() {
                public X509Certificate[] getAcceptedIssuers() { return new X509Certificate[0]; }
                public void checkClientTrusted(X509Certificate[] certs, String authType) {}
                public void checkServerTrusted(X509Certificate[] certs, String authType) {}
            }
        };

        ctx.init(null, trustAllCerts, new SecureRandom());
        return ctx;
    }

    @Test
    public void testWrapUnwrapBasic() throws Exception {
        SSLContext ctx = createTestSSLContext();

        SSLEngine clientEngine = ctx.createSSLEngine("localhost", 443);
        clientEngine.setUseClientMode(true);

        SSLEngine serverEngine = ctx.createSSLEngine();
        serverEngine.setUseClientMode(false);

        // Begin handshake on both sides
        clientEngine.beginHandshake();
        serverEngine.beginHandshake();

        // Allocate buffers
        ByteBuffer clientAppOut = ByteBuffer.allocate(APP_BUFFER_SIZE);
        ByteBuffer clientNetOut = ByteBuffer.allocate(NET_BUFFER_SIZE);
        ByteBuffer serverAppOut = ByteBuffer.allocate(APP_BUFFER_SIZE);
        ByteBuffer serverNetOut = ByteBuffer.allocate(NET_BUFFER_SIZE);

        // Client should be in NEED_WRAP state
        assertEquals("Client should need to wrap",
                     SSLEngineResult.HandshakeStatus.NEED_WRAP,
                     clientEngine.getHandshakeStatus());

        // Server should be in NEED_UNWRAP state (waiting for ClientHello)
        assertEquals("Server should need to unwrap",
                     SSLEngineResult.HandshakeStatus.NEED_UNWRAP,
                     serverEngine.getHandshakeStatus());

        // Test that wrap produces output
        SSLEngineResult result = clientEngine.wrap(clientAppOut, clientNetOut);
        assertNotNull("Wrap result should not be null", result);
        assertTrue("Wrap should produce bytes", result.bytesProduced() > 0 ||
                   result.getHandshakeStatus() != SSLEngineResult.HandshakeStatus.NOT_HANDSHAKING);
    }

    @Test
    public void testWrapWithEmptySource() throws Exception {
        SSLContext ctx = createTestSSLContext();

        SSLEngine engine = ctx.createSSLEngine("localhost", 443);
        engine.setUseClientMode(true);
        engine.beginHandshake();

        // Empty source buffer
        ByteBuffer emptySource = ByteBuffer.allocate(0);
        ByteBuffer netOut = ByteBuffer.allocate(NET_BUFFER_SIZE);

        // Wrap with empty source should still work (produces handshake data)
        SSLEngineResult result = engine.wrap(emptySource, netOut);
        assertNotNull("Result should not be null", result);
        assertEquals("Status should be OK", SSLEngineResult.Status.OK, result.getStatus());
    }

    @Test
    public void testUnwrapWithEmptySource() throws Exception {
        SSLContext ctx = createTestSSLContext();

        SSLEngine engine = ctx.createSSLEngine();
        engine.setUseClientMode(false); // Server mode
        engine.beginHandshake();

        // Empty source buffer
        ByteBuffer emptySource = ByteBuffer.allocate(0);
        ByteBuffer appOut = ByteBuffer.allocate(APP_BUFFER_SIZE);

        // Unwrap with empty source
        SSLEngineResult result = engine.unwrap(emptySource, appOut);
        assertNotNull("Result should not be null", result);
        // Should return OK or BUFFER_UNDERFLOW since no data is available
        assertTrue("Status should be OK or BUFFER_UNDERFLOW",
                   result.getStatus() == SSLEngineResult.Status.OK ||
                   result.getStatus() == SSLEngineResult.Status.BUFFER_UNDERFLOW);
    }

    @Test(expected = IllegalArgumentException.class)
    public void testWrapWithNullSource() throws Exception {
        SSLContext ctx = createTestSSLContext();

        SSLEngine engine = ctx.createSSLEngine("localhost", 443);
        engine.setUseClientMode(true);
        engine.beginHandshake();

        ByteBuffer netOut = ByteBuffer.allocate(NET_BUFFER_SIZE);

        // Wrap with null source should throw
        engine.wrap((ByteBuffer[]) null, netOut);
    }

    @Test(expected = IllegalArgumentException.class)
    public void testWrapWithNullDestination() throws Exception {
        SSLContext ctx = createTestSSLContext();

        SSLEngine engine = ctx.createSSLEngine("localhost", 443);
        engine.setUseClientMode(true);
        engine.beginHandshake();

        ByteBuffer appOut = ByteBuffer.allocate(APP_BUFFER_SIZE);

        // Wrap with null destination should throw
        engine.wrap(appOut, null);
    }

    @Test(expected = IllegalArgumentException.class)
    public void testUnwrapWithNullSource() throws Exception {
        SSLContext ctx = createTestSSLContext();

        SSLEngine engine = ctx.createSSLEngine();
        engine.setUseClientMode(false);
        engine.beginHandshake();

        ByteBuffer appOut = ByteBuffer.allocate(APP_BUFFER_SIZE);

        // Unwrap with null source should throw
        engine.unwrap(null, appOut);
    }

    @Test
    public void testBufferOverflow() throws Exception {
        SSLContext ctx = createTestSSLContext();

        SSLEngine engine = ctx.createSSLEngine("localhost", 443);
        engine.setUseClientMode(true);
        engine.beginHandshake();

        ByteBuffer appOut = ByteBuffer.allocate(APP_BUFFER_SIZE);
        // Very small destination buffer
        ByteBuffer smallNetOut = ByteBuffer.allocate(10);

        SSLEngineResult result = engine.wrap(appOut, smallNetOut);
        // Should indicate BUFFER_OVERFLOW if output doesn't fit
        // Note: This depends on implementation - some may return OK with 0 bytes
        assertNotNull("Result should not be null", result);
    }

    @Test
    public void testMultipleWrapCalls() throws Exception {
        SSLContext ctx = createTestSSLContext();

        SSLEngine engine = ctx.createSSLEngine("localhost", 443);
        engine.setUseClientMode(true);
        engine.beginHandshake();

        ByteBuffer appOut = ByteBuffer.allocate(APP_BUFFER_SIZE);
        ByteBuffer netOut = ByteBuffer.allocate(NET_BUFFER_SIZE);

        // First wrap
        SSLEngineResult result1 = engine.wrap(appOut, netOut);
        assertNotNull("First wrap result should not be null", result1);

        // Prepare for second wrap
        netOut.flip();
        netOut.compact();

        // Second wrap (may need to continue handshake)
        SSLEngineResult result2 = engine.wrap(appOut, netOut);
        assertNotNull("Second wrap result should not be null", result2);
    }

    @Test
    public void testHandshakeStatusTransitions() throws Exception {
        SSLContext ctx = createTestSSLContext();

        SSLEngine clientEngine = ctx.createSSLEngine("localhost", 443);
        clientEngine.setUseClientMode(true);

        // Before beginHandshake
        assertEquals("Initial status should be NOT_HANDSHAKING",
                     SSLEngineResult.HandshakeStatus.NOT_HANDSHAKING,
                     clientEngine.getHandshakeStatus());

        // After beginHandshake
        clientEngine.beginHandshake();
        assertEquals("After begin, client should NEED_WRAP",
                     SSLEngineResult.HandshakeStatus.NEED_WRAP,
                     clientEngine.getHandshakeStatus());
    }

    @Test
    public void testCloseAfterHandshakeStart() throws Exception {
        SSLContext ctx = createTestSSLContext();

        SSLEngine engine = ctx.createSSLEngine("localhost", 443);
        engine.setUseClientMode(true);
        engine.beginHandshake();

        // Close outbound during handshake
        engine.closeOutbound();
        assertTrue("Outbound should be done", engine.isOutboundDone());

        // Engine should still be able to close inbound
        engine.closeInbound();
        assertTrue("Inbound should be done", engine.isInboundDone());
    }

    @Test
    public void testReadOnlyBufferException() throws Exception {
        SSLContext ctx = createTestSSLContext();

        SSLEngine engine = ctx.createSSLEngine("localhost", 443);
        engine.setUseClientMode(true);
        engine.beginHandshake();

        ByteBuffer appOut = ByteBuffer.allocate(APP_BUFFER_SIZE);
        ByteBuffer netOut = ByteBuffer.allocate(NET_BUFFER_SIZE).asReadOnlyBuffer();

        try {
            engine.wrap(appOut, netOut);
            fail("Should throw exception for read-only destination buffer");
        } catch (java.nio.ReadOnlyBufferException e) {
            // Expected
        }
    }

    @Test
    public void testWrapArrayOfBuffers() throws Exception {
        SSLContext ctx = createTestSSLContext();

        SSLEngine engine = ctx.createSSLEngine("localhost", 443);
        engine.setUseClientMode(true);
        engine.beginHandshake();

        // Create array of source buffers
        ByteBuffer[] srcs = new ByteBuffer[] {
            ByteBuffer.allocate(100),
            ByteBuffer.allocate(100),
            ByteBuffer.allocate(100)
        };
        ByteBuffer dst = ByteBuffer.allocate(NET_BUFFER_SIZE);

        SSLEngineResult result = engine.wrap(srcs, dst);
        assertNotNull("Result should not be null", result);
    }

    @Test
    public void testUnwrapArrayOfBuffers() throws Exception {
        SSLContext ctx = createTestSSLContext();

        SSLEngine engine = ctx.createSSLEngine();
        engine.setUseClientMode(false);
        engine.beginHandshake();

        ByteBuffer src = ByteBuffer.allocate(NET_BUFFER_SIZE);
        ByteBuffer[] dsts = new ByteBuffer[] {
            ByteBuffer.allocate(100),
            ByteBuffer.allocate(100),
            ByteBuffer.allocate(100)
        };

        SSLEngineResult result = engine.unwrap(src, dsts);
        assertNotNull("Result should not be null", result);
    }

    @Test
    public void testSessionCreation() throws Exception {
        SSLContext ctx = createTestSSLContext();

        SSLEngine engine = ctx.createSSLEngine("localhost", 443);
        engine.setUseClientMode(true);

        // Session creation should be enabled by default
        assertTrue("Session creation should be enabled", engine.getEnableSessionCreation());

        // Test setting it
        engine.setEnableSessionCreation(false);
        // Note: HiTLS implementation always returns true currently
    }
}
